import os
import numpy as np
from PIL import ImageEnhance, Image, ImageFilter, ImageOps
import random
import uuid
import math

# 配置常量
ORIGINAL_VARIATIONS = 10  # 原始增强方法的变体数量
NEW_VARIATIONS = 5  # 新增增强方法的变体数量
BRIGHTNESS_RANGE = (0.5, 1.5)
CONTRAST_RANGE = (0.5, 1.5)
ROTATION_RANGE = (0, 180)
BLUR_RANGE = (0.5, 2.0)
COLOR_RANGE = (0.5, 1.5)
SATURATION_RANGE = (0.5, 1.5)
SHARPNESS_RANGE = (0.5, 1.5)
EDGE_RANGE = (0.5, 2.0)
CROP_RANGE = (0.5, 0.9)
JITTER_RANGE = (0.1, 0.3)
PERSPECTIVE_RANGE = (0.1, 0.4)
NOISE_RANGE = (0.05, 0.2)


# 1. 色度增强 - 5个变体
def color_enhancement(image, factors):
    """生成多个色度增强版本"""
    return [ImageEnhance.Color(image).enhance(f).convert('RGB') for f in factors]


# 2. 锐度增强 - 5个变体
def sharpness_enhancement(image, factors):
    """生成多个锐度增强版本"""
    return [ImageEnhance.Sharpness(image).enhance(f).convert('RGB') for f in factors]


# 3. 高斯模糊 - 5个变体
def gaussian_blur(image, radii):
    """生成多个模糊版本"""
    return [image.filter(ImageFilter.GaussianBlur(r)).convert('RGB') for r in radii]


# 4. 边缘增强 - 5个变体
def edge_enhancement(image, strengths):
    """生成多个边缘增强版本"""
    results = []
    for s in strengths:
        # 使用边缘增强滤波器
        enhanced = image.filter(ImageFilter.EDGE_ENHANCE_MORE)
        # 额外应用锐化以增强效果
        results.append(ImageEnhance.Sharpness(enhanced).enhance(s).convert('RGB'))
    return results


# 5. 随机裁剪 - 5个变体
def random_crop(image, crop_percentages):
    """生成多个裁剪版本"""
    results = []
    width, height = image.size

    for p in crop_percentages:
        # 基于比例计算裁剪尺寸
        crop_width = int(width * p)
        crop_height = int(height * p)

        # 随机位置裁剪
        left = random.randint(0, width - crop_width)
        top = random.randint(0, height - crop_height)
        right = left + crop_width
        bottom = top + crop_height

        # 裁剪并恢复原始尺寸
        cropped = image.crop((left, top, right, bottom))
        results.append(cropped.resize((width, height), Image.LANCZOS).convert('RGB'))
    return results


# 6. 饱和度调整 - 5个变体
def saturation_adjustment(image, factors):
    """生成多个饱和度调整版本"""
    results = []
    for f in factors:
        if f < 1.0:
            # 降低饱和度
            gray = image.convert('L').convert('RGB')
            result = Image.blend(image, gray, 1.0 - f)
        else:
            # 增加饱和度
            result = ImageEnhance.Color(image).enhance(f)
        results.append(result.convert('RGB'))
    return results


# 7. 直方图均衡化 - 单一效果
def histogram_equalization(image):
    """应用直方图均衡化"""
    return ImageOps.equalize(image).convert('RGB')


# 8. 随机抖动 - 5个变体
def color_jitter(image, jitter_strengths):
    """生成多个颜色抖动版本"""
    results = []
    for s in jitter_strengths:
        # 分离通道
        r, g, b = image.split()

        # 为每个通道应用不同的增强因子
        factors = [1 + random.uniform(-s, s) for _ in range(3)]

        # 应用增强并重新合并
        r = ImageEnhance.Brightness(r).enhance(factors[0])
        g = ImageEnhance.Brightness(g).enhance(factors[1])
        b = ImageEnhance.Brightness(b).enhance(factors[2])
        results.append(Image.merge('RGB', (r, g, b)).convert('RGB'))
    return results


# 9. 透视变换 - 5个变体
def perspective_transform(image, distortions):
    """生成多个透视变换版本"""
    results = []
    width, height = image.size
    for d in distortions:
        # 随机位移
        dx1 = random.uniform(-d, d) * width
        dy1 = random.uniform(-d, d) * height
        dx2 = random.uniform(-d, d) * width
        dy2 = random.uniform(-d, d) * height

        # 定义四个角的原始坐标和目标坐标
        original_corners = [(0, 0), (width, 0), (width, height), (0, height)]
        target_corners = [
            (dx1, dy1),
            (width + dx2, dy1),
            (width - dx1, height + dy2),
            (dx2, height - dy2)
        ]

        # 计算变换矩阵
        coefficients = []
        for (x1, y1), (x2, y2) in zip(original_corners, target_corners):
            coefficients.extend([
                [x1, y1, 1, 0, 0, 0, -x1 * x2, -y1 * x2],
                [0, 0, 0, x1, y1, 1, -x1 * y2, -y1 * y2]
            ])

        # 计算变换参数
        A = np.array(coefficients)
        B = np.array([(x2, y2) for (x2, y2) in target_corners]).flatten()

        # 解决最小二乘问题
        try:
            perspective_coeffs = np.linalg.lstsq(A, B, rcond=None)[0]
            perspective_matrix = list(perspective_coeffs[:8]) + [1.0]
        except:
            perspective_matrix = (1, 0, 0, 0, 1, 0, 0, 0)

        # 应用透视变换
        transformed = image.transform(
            (width, height),
            Image.PERSPECTIVE,
            perspective_matrix,
            Image.BICUBIC,
            fillcolor=(255, 255, 255)
        )
        results.append(transformed.convert('RGB'))
    return results


# 10. 随机噪声 - 5个变体
def add_noise(image, noise_levels):
    """生成多个噪声版本"""
    results = []
    img_array = np.array(image).astype(np.float32) / 255.0
    for level in noise_levels:
        # 添加高斯噪声
        noise = np.random.normal(0, level, img_array.shape)
        noisy_img = np.clip(img_array + noise, 0, 1.0) * 255
        results.append(Image.fromarray(noisy_img.astype(np.uint8)).convert('RGB'))
    return results


# 11. 反相颜色 - 单一效果
def invert_colors(image):
    """反相图像颜色"""
    return ImageOps.invert(image).convert('RGB')


# 12. 海报化 - 5个变体
def posterize(image, bits_values):
    """生成多个海报化版本"""
    return [ImageOps.posterize(image, b).convert('RGB') for b in bits_values]


# 保留原始增强方法
def brightnessEnhancement(image, factors):
    """生成多个亮度增强版本"""
    return [ImageEnhance.Brightness(image).enhance(f).convert('RGB') for f in factors]


def contrastEnhancement(image, factors):
    """生成多个对比度增强版本"""
    return [ImageEnhance.Contrast(image).enhance(f).convert('RGB') for f in factors]


def rotation(image, angles):
    """生成多个旋转版本"""
    return [
        image.rotate(angle, expand=True, fillcolor=(255, 255, 255)).convert('RGB')
        for angle in angles
    ]


def flip(image):
    """水平翻转"""
    return image.transpose(Image.FLIP_LEFT_RIGHT).convert('RGB')


def save_image(image, save_dir, base_name, op_name, suffix):
    """保存图像并确保文件名唯一"""
    while True:
        rand_part = uuid.uuid4().hex[:8]
        filename = f"{base_name}_{op_name}_{suffix}_{rand_part}.png"
        full_path = os.path.join(save_dir, filename)
        if not os.path.exists(full_path):
            image.save(full_path)
            return


def createImage(imageDir, saveDir):
    """应用所有增强技术"""
    if not os.path.exists(saveDir):
        os.makedirs(saveDir)

    # 预计算参数
    brightness_factors = np.linspace(*BRIGHTNESS_RANGE, ORIGINAL_VARIATIONS)
    contrast_factors = np.linspace(*CONTRAST_RANGE, ORIGINAL_VARIATIONS)
    rotation_angles = np.linspace(*ROTATION_RANGE, ORIGINAL_VARIATIONS)

    # 新增方法的参数
    color_factors = np.linspace(*COLOR_RANGE, NEW_VARIATIONS)
    sharpness_factors = np.linspace(*SHARPNESS_RANGE, NEW_VARIATIONS)
    blur_radii = np.linspace(*BLUR_RANGE, NEW_VARIATIONS)
    edge_strengths = np.linspace(*EDGE_RANGE, NEW_VARIATIONS)
    crop_percentages = np.linspace(*CROP_RANGE, NEW_VARIATIONS)
    saturation_factors = np.linspace(*SATURATION_RANGE, NEW_VARIATIONS)
    jitter_strengths = np.linspace(*JITTER_RANGE, NEW_VARIATIONS)
    perspective_distortions = np.linspace(*PERSPECTIVE_RANGE, NEW_VARIATIONS)
    noise_levels = np.linspace(*NOISE_RANGE, NEW_VARIATIONS)
    posterize_bits = [random.randint(3, 6) for _ in range(NEW_VARIATIONS)]

    # 过滤掉零度旋转避免重复
    rotation_angles = [a for a in rotation_angles if abs(a) > 1e-3]

    for img_name in os.listdir(imageDir):
        img_path = os.path.join(imageDir, img_name)
        if not os.path.isfile(img_path):
            continue

        try:
            with Image.open(img_path) as img:
                # 转换为RGB并缓存
                img = img.convert('RGB')
                base_name = os.path.splitext(img_name)[0]

                # 原始操作
                # 1. 亮度增强 - 10个变体
                bright_imgs = brightnessEnhancement(img, brightness_factors)
                for i, bright_img in enumerate(bright_imgs):
                    save_image(bright_img, saveDir, base_name, 'bright', i)

                # 2. 对比度增强 - 10个变体
                contrast_imgs = contrastEnhancement(img, contrast_factors)
                for i, contrast_img in enumerate(contrast_imgs):
                    save_image(contrast_img, saveDir, base_name, 'contrast', i)

                # 3. 旋转 - 10个变体
                rotated_imgs = rotation(img, rotation_angles)
                for i, rotated_img in enumerate(rotated_imgs):
                    save_image(rotated_img, saveDir, base_name, 'rotate', i)

                # 4. 翻转 - 单一效果
                flipped_img = flip(img)
                save_image(flipped_img, saveDir, base_name, 'flip', 0)

                # 新增操作
                # 5. 色度增强 - 5个变体
                color_imgs = color_enhancement(img, color_factors)
                for i, color_img in enumerate(color_imgs):
                    save_image(color_img, saveDir, base_name, 'color', i)

                # 6. 锐度增强 - 5个变体
                sharp_imgs = sharpness_enhancement(img, sharpness_factors)
                for i, sharp_img in enumerate(sharp_imgs):
                    save_image(sharp_img, saveDir, base_name, 'sharpness', i)

                # 7. 高斯模糊 - 5个变体
                blur_imgs = gaussian_blur(img, blur_radii)
                for i, blur_img in enumerate(blur_imgs):
                    save_image(blur_img, saveDir, base_name, 'blur', i)

                # 8. 边缘增强 - 5个变体
                edge_imgs = edge_enhancement(img, edge_strengths)
                for i, edge_img in enumerate(edge_imgs):
                    save_image(edge_img, saveDir, base_name, 'edge', i)

                # 9. 随机裁剪 - 5个变体
                crop_imgs = random_crop(img, crop_percentages)
                for i, crop_img in enumerate(crop_imgs):
                    save_image(crop_img, saveDir, base_name, 'crop', i)

                # 10. 饱和度调整 - 5个变体
                sat_imgs = saturation_adjustment(img, saturation_factors)
                for i, sat_img in enumerate(sat_imgs):
                    save_image(sat_img, saveDir, base_name, 'saturation', i)

                # 11. 直方图均衡化 - 单一效果
                eq_img = histogram_equalization(img)
                save_image(eq_img, saveDir, base_name, 'equalize', 0)

                # 12. 随机抖动 - 5个变体
                jitter_imgs = color_jitter(img, jitter_strengths)
                for i, jitter_img in enumerate(jitter_imgs):
                    save_image(jitter_img, saveDir, base_name, 'jitter', i)

                # 13. 透视变换 - 5个变体
                persp_imgs = perspective_transform(img, perspective_distortions)
                for i, persp_img in enumerate(persp_imgs):
                    save_image(persp_img, saveDir, base_name, 'perspective', i)

                # 14. 随机噪声 - 5个变体
                noisy_imgs = add_noise(img, noise_levels)
                for i, noisy_img in enumerate(noisy_imgs):
                    save_image(noisy_img, saveDir, base_name, 'noise', i)

                # 15. 反相颜色 - 单一效果
                invert_img = invert_colors(img)
                save_image(invert_img, saveDir, base_name, 'invert', 0)

                # 16. 海报化 - 5个变体
                poster_imgs = posterize(img, posterize_bits)
                for i, poster_img in enumerate(poster_imgs):
                    save_image(poster_img, saveDir, base_name, 'posterize', i)

        except Exception as e:
            print(f"处理图像 {img_name} 时出错: {str(e)}")


if __name__ == '__main__':
    # 使用相对路径
    imageDir = "images"
    saveDir = "images_enhancement"

    # 创建输出目录
    if not os.path.exists(saveDir):
        os.makedirs(saveDir)

    createImage(imageDir, saveDir)
    print(f"图像增强完成! 共生成 {ORIGINAL_VARIATIONS * 3 + 1 + NEW_VARIATIONS * 10 + 3} 个变体/原始图像")
    print(f"输出保存在: {saveDir}")